#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>

void GemmFloatUnit_4(float *dst, float *input, float *weight, size_t src_depth, size_t dst_step, size_t dst_depth, size_t weight_offset);
void GemmFloatUnit_4_old(float *dst, float *input, float *weight, size_t src_depth, size_t dst_step, size_t dst_depth, size_t weight_offset);

void main()
{
    int iw = 32;
    int ih = 32;
    int ic = 8;
    int oc = 8;
    int kw = 3;
    int kh = 3;
    int ow = 32;
    int oh = 32;

    int packed_input_size = ow * oh * kw * kh * ic;
    float *packed_input_1 = (float *) malloc(packed_input_size * sizeof(float));
    float *packed_input_2 = (float *) malloc(packed_input_size * sizeof(float));
    for (int i = 0; i < packed_input_size; ++i) {
        packed_input_1[i] = (float) (i % 10);
        packed_input_2[i] = packed_input_1[i];
    }

    int packed_weight_size = oc * kw * kh * ic;
    float *packed_weight_1 = (float *) malloc(packed_weight_size * sizeof(float));
    float *packed_weight_2 = (float *) malloc(packed_weight_size * sizeof(float));
    for (int i = 0; i < packed_weight_size; ++i) {
        packed_weight_1[i] = (float) (i % 5 + 2);
        packed_weight_2[i] = packed_weight_1[i];
    }

    int out_size = ow * oh * oc;
    int *output = (int *) malloc(out_size * sizeof(int));
    memset(output, 0, out_size * sizeof(int));

    struct timeval start, end;
    double duration = 0;
    int batches = ow * oh * oc / 8 / 8;
    gettimeofday(&start, NULL);
    for (int i  = 0; i < batches; i++) {
        GemmFloatUnit_4_old(output, packed_input + i * kh * kw * 4, packed_weight, kw * kh * ic / 4, ow * oh * 4, oc / 4, 0);
    }
    gettimeofday(&end, NULL);
    duration = (end.tv_sec - start.tv_sec) * 1000000 + (end.tv_usec - start.tv_usec);
    printf("%f\n", duration);
    gettimeofday(&start, NULL);
    for (int i  = 0; i < batches; i++) {
        GemmFloatUnit_4(output, packed_input + i * kh * kw * 4, packed_weight, kw * kh * ic / 4, ow * oh * 4, oc / 4, 0);
    }
    gettimeofday(&end, NULL);
    duration = (end.tv_sec - start.tv_sec) * 1000000 + (end.tv_usec - start.tv_usec);
    printf("%f\n", duration);
}
